TYPE=$1
CPU_IMAGE=tensorflow/tensorflow:2.6.0-jupyter
GPU_IMAGE=tensorflow/tensorflow:2.6.0-gpu-jupyter
TARGET_REPO=registry.cn-hangzhou.aliyuncs.com/all3n/finenv
if [[ "$TYPE" == "cpu" ]];then
    IMAGE=$CPU_IMAGE
elif [[ "$TYPE" == "gpu" ]];then
    IMAGE=$GPU_IMAGE
else
    IMAGE=all
fi

: ${PUSH:=false}

build_image(){
    echo $1 $2 $PUSH
    docker build --build-arg BASE=$1 -t $TARGET_REPO:$2 .
    if [[ "$PUSH" == "true" ]];then
        docker push $TARGET_REPO:$2
    fi
}


if [[ "$IMAGE" == "all" ]];then
    build_image $CPU_IMAGE cpu
    build_image $GPU_IMAGE gpu
else
    build_image $IMAGE $TYPE
fi
